import os
import sys

import torch

from .others import Timer


def fit(
        epoches, data_loader, model, loss_func,
        optimzer=None, device=None,
        lr_scheduler=None,
        save_path_folder=None,
        save_model_epoches=None,
        print_info_batches=100
        ):
    fit_loss, fit_acc = [], []
    if not device:
        device = torch.device('cpu')
    if not optimzer:
        optimzer = torch.optim.Adam(model.parameters())
    save_flag = save_path_folder and save_model_epoches
    if not save_flag and (save_path_folder or save_model_epoches):
        print(
            'you should give both "save_path_folder" and "save_model_epoches"'
            'if you want to save the training model.', file=sys.stderr
            )
    timer = Timer()
    model.train()
    with Timer("fit the model"):
        for epoch in range(1, epoches + 1):
            loss, acc = train(
                    epoch, data_loader, model, loss_func, optimzer, device,
                    print_info_batches, timer
                    )
            if not lr_scheduler:
                lr_scheduler.step()
            fit_acc.append(acc)
            fit_loss.append(loss)
            if save_flag and epoch % save_model_epoches == 0:
                torch.save(
                        model.state_dict(),
                        os.path.join(save_path_folder, '{:04d}.pt'.format(epoch))
                        )
    return fit_loss, fit_acc


def train(
        epoch, loader, model, loss_func, optimzer, device=torch.device("cpu"),
        print_info_batches=100,
        timer=None
        ) -> tuple:
    run_loss = 0
    total = 0
    correct = 0
    if not timer:
        timer = Timer()
    timer.start()
    for batch_idx, (X, Y) in enumerate(loader):
        X, Y = X.to(device), Y.to(device)

        Y_pred = model(X)

        loss = loss_func(Y_pred, Y)

        optimzer.zero_grad()
        loss.backward()
        optimzer.step()

        run_loss += loss.item()
        total += Y.size(0)
        # 使用max函数获取指定维度上的最大值的值与索引（下标），返回值第一个是值，第二个是索引。
        pred_y = Y_pred.detach().argmax(dim=1)
        correct += (pred_y.detach() == Y.detach()).sum().item()
        if batch_idx % print_info_batches + 1 == print_info_batches:
            print('batch = ', batch_idx + 1, 'sum loss = ', run_loss, sep=None)
    print(
            f'Train {epoch}: total = {total}, loss = {run_loss / total}, acc = '
            f'{correct / total}, using time = {timer.stop(False)}s'
            )
    return run_loss / total, correct / total
